package qq2564874169.tuaotuao.fx.mq.zookeeper;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import qq2564874169.tuaotuao.fx.*;
import qq2564874169.tuaotuao.fx.mq.MqMessage;
import qq2564874169.tuaotuao.fx.mq.MqSubscriber;
import qq2564874169.tuaotuao.fx.mq.SubscriberError;
import qq2564874169.tuaotuao.fx.utils.ExUtil;

import java.nio.charset.StandardCharsets;
import java.util.*;

public class ZkMqSubscriber implements MqSubscriber {
    public static final String msgNode = "msg";
    private static String rootNode;
    private ZooKeeper zk;
    private static String receivedKey = UUID.randomUUID().toString();
    private static String errorKey = UUID.randomUUID().toString();
    private EventDispatcher<MqMessage> receivedDispatcher = new EventDispatcher<>(receivedKey);
    private EventDispatcher<SubscriberError> errorDispatcher = new EventDispatcher<>(errorKey);
    private Map<String, MsgWatch> consumeMap = new HashMap<>(8);

    public static void createRootNode(ZooKeeper zk) {
        createRootNode(zk, "/FxMsgQueue");
    }

    public static void createRootNode(ZooKeeper zk, String nodeName) {
        try {
            if (zk.exists(nodeName, false) == null) {
                zk.create(nodeName, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
            rootNode = nodeName;
        } catch (KeeperException | InterruptedException e) {
            throw ExUtil.toRuntime(e);
        }
    }

    private static String fullPath(String path) {
        FxAssert.isNotNull(path);
        FxAssert.isNotNull(rootNode, "请先调用createRootNode创建根节点");
        if (path.startsWith("/")) {
            return rootNode + path;
        } else {
            return rootNode + "/" + path;
        }
    }

    public ZkMqSubscriber(ZooKeeper zk) {
        this.zk = zk;
    }

    @Override
    public void declare(String queue) {
        var k = fullPath(queue);
        try {
            if (zk.exists(k, false) == null) {
                zk.create(k, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (KeeperException | InterruptedException e) {
            throw ExUtil.toRuntime(e);
        }
    }

    @Override
    public EventDispatcher<MqMessage> evReceived() {
        return receivedDispatcher;
    }

    @Override
    public EventDispatcher<SubscriberError> evError() {
        return errorDispatcher;
    }

    @Override
    public String consume(String queue, String name, Object state) {
        var qkey = fullPath(queue);
        var nkey = qkey + "/" + name;
        try {
            FxAssert.isNotNull(zk.exists(qkey, false), queue + "不存在");
            var nstat = zk.exists(nkey, false);
            if (nstat == null) {
                try {
                    zk.create(nkey, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                } catch (KeeperException e) {
                    if (e.code().equals(KeeperException.Code.NODEEXISTS) == false) {
                        throw ExUtil.toRuntime(e);
                    }
                }
            }
            var msgWatch = new MsgWatch(queue, name, state);
            Task.run(() -> {
                try {
                    var childs = zk.getChildren(nkey, false);
                    for (String c : childs) {
                        msgWatch.process(new WatchedEvent(
                                Watcher.Event.EventType.NodeChildrenChanged,
                                Watcher.Event.KeeperState.Expired,
                                nkey + "/" + c));
                    }
                    zk.addWatch(nkey, msgWatch, AddWatchMode.PERSISTENT);
                } catch (KeeperException | InterruptedException e) {
                    SubscriberError exArgs = new SubscriberError(e, state);
                    errorDispatcher.fire(errorKey, ZkMqSubscriber.this, exArgs);
                    if (exArgs.isRetry) {
                        consume(queue, name, state);
                    }
                }
            });
            consumeMap.put(msgWatch.id, msgWatch);
            return msgWatch.id;
        } catch (Exception ex) {
            SubscriberError exArgs = new SubscriberError(ex, state);
            errorDispatcher.fire(errorKey, ZkMqSubscriber.this, exArgs);
            if (exArgs.isRetry) {
                return consume(queue, name, state);
            } else {
                throw ExUtil.toRuntime(ex);
            }
        }
    }

    @Override
    public void cancel(String queue, String name) {
        FxStream.create(consumeMap.values())
                .filter(i -> i.queue.equals(queue) && i.name.equals(name))
                .foreach(i -> consumeMap.remove(i.id).stop());
    }

    @Override
    public void cancel(String consumeId) {
        var w = consumeMap.remove(consumeId);
        if (w != null) {
            w.stop();
        }
    }

    @Override
    public void close() {
        FxStream.create(consumeMap.values()).foreach(i -> i.stop());
        consumeMap.clear();
    }

    private class MsgWatch implements Watcher {
        public final String name;
        public final String queue;
        public final String fullkey;
        public final Object state;
        public final String id;
        private boolean isStop;

        public MsgWatch(String queue, String name, Object state) {
            this.id = UUID.randomUUID().toString();
            this.queue = queue;
            this.name = name;
            this.state = state;
            this.fullkey = fullPath(queue) + "/" + name;
        }

        public void stop() {
            isStop = true;
        }

        @Override
        public void process(WatchedEvent event) {
            if (isStop) {
                return;
            }
            if (event.getState().equals(Watcher.Event.KeeperState.Closed)) {
                return;
            }
            if (!event.getType().equals(Event.EventType.NodeChildrenChanged)) {
                return;
            }
            try {
                var childs = FxStream.create(zk.getChildren(fullkey, false));
                if (!childs.any()) {
                    return;
                }
                childs = childs.filter(i -> i.startsWith(msgNode) && i.length() > msgNode.length())
                        .sort(i -> Long.parseLong(i.substring(msgNode.length() + 1)));
                var msgkey = new Var<String>();
                var lockkey = new Var<String>();

                while (childs.any()) {
                    try {
                        var name = childs.first().orElse(null);
                        msgkey.set(fullkey + "/" + name);
                        lockkey.set(msgkey.get() + "/lock");
                        if (zk.exists(lockkey.get(), false) != null) {
                            throw KeeperException.create(KeeperException.Code.NODEEXISTS);
                        } else {
                            zk.create(lockkey.get(), new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
                        }
                    } catch (KeeperException e) {
                        if (e.code().equals(KeeperException.Code.NODEEXISTS)) {
                            childs.removeRange(0, 1);
                        } else {
                            throw e;
                        }
                    }
                }
                if (!childs.any()) {
                    return;
                }
                var data = zk.getData(msgkey.get(), false, new Stat());
                if (data.length < 1) {
                    try {
                        zk.delete(lockkey.get(), -1);
                    } catch (KeeperException e) {
                        if (!e.code().equals(KeeperException.Code.NONODE)) {
                            throw e;
                        }
                    }
                    try {
                        zk.delete(msgkey.get(), -1);
                    } catch (KeeperException e) {
                        if (!e.code().equals(KeeperException.Code.NONODE)) {
                            throw e;
                        }
                    }
                    return;
                }
                MqMessage msg = new MqMessage();
                msg.queue = queue;
                msg.name = name;
                msg.state = state;
                msg.data = new String(data, StandardCharsets.UTF_8);
                try {
                    receivedDispatcher.fire(receivedKey, ZkMqSubscriber.this, msg);
                } catch (Exception ex) {
                    consumeMap.remove(fullkey);
                    throw ex;
                }
                if (msg.reconsume) {
                    zk.transaction()
                            .delete(lockkey.get(), -1).delete(msgkey.get(), -1)
                            .create(msgkey.get(), data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT)
                            .commit();
                } else {
                    zk.transaction()
                            .delete(lockkey.get(), -1).delete(msgkey.get(), -1)
                            .commit();
                }
            } catch (KeeperException | InterruptedException e) {
                if (zk.getState().equals(ZooKeeper.States.CLOSED) == false) {
                    throw ExUtil.toRuntime(e);
                }
            }
        }
    }
}
